import requests
import json

# Ollama API 地址
OLLAMA_API_URL = "http://localhost:11434"
MODEL_NAME = "deepseek-r1:1.5b"

# 系统提示
SYSTEM = """
你是一个严谨的AI助手，必须用中文回答问题。
你的回答必须清晰，准确，并且遵循以下格式：
User: [用户输入]
Assistant: [你的回答]
如果用户的问题不明确，请要求用户提供更多信息。
"""

# 模板
def build_prompt(system_prompt, user_input):
    return f"""<|system|>
{system_prompt}<|end|>
<|user|>
User: {user_input}<|end|>
<|assistant|>
Assistant:"""



continue_flag = True  # 控制循环

while continue_flag:

    # 手动输入用户问题
    user_input = input("请输入你的问题：")  # 手动输入用户问题

    if user_input == "退出":
        continue_flag = False  # 退出循环
        print("感谢使用！")
        break

    # 组织请求数据
    data = {
        "model": MODEL_NAME,
        "prompt": build_prompt(SYSTEM, user_input),  # 使用模板构建完整的提示
    }


    try:
        # 发送 POST 请求
        # print("正在发送请求...")
        response = requests.post(f"{OLLAMA_API_URL}/api/generate", json=data, timeout=30, stream=True)
        # print("请求完成，状态码：", response.status_code)

        # 打印 API 响应
        if response.status_code == 200:
            # print("模型回复：", response.text)  # 直接打印原始响应

            full_response = "" # 完整回复

            for line in response.iter_lines(): 
                if line:
                    # 解析 JSON
                    chunk = line.decode("utf-8")  # 将字节流转换为字符串
                    json_data = json.loads(chunk)  # 将字符串转换为字典

                    if "response" in json_data:
                        full_response += json_data["response"]
                    if json_data.get("done", False):
                        break  # 如果 done 为 True，结束循环
            

            if "<think>" in full_response and "</think>" in full_response:
                start = full_response.find("<think>")
                end = full_response.find("</think>")
                full_response = full_response[:start] + full_response[end + len("</think>"):]
                full_response = full_response.strip()


            print("DeepSeek回复：", full_response)


        else:
            print(f"请求失败，状态码：{response.status_code}")
            print("API 返回：", response.text)
    except requests.exceptions.Timeout:
        print("请求超时，请检查网络或 API 服务是否正常。")
    except requests.exceptions.RequestException as e:
        print(f"请求发生错误：{e}")
